# !/usr/bin/env python3
# coding=utf8
"""
程序载入"json配置文件",
并读取 vnpy.tools.database_cta.CTA_BACKTEST_RESULT_BF_CONT_DBS 目录下的 sqlite 文件,
计算之后写入"database_cta_cont.db"文件,
"""
import argparse
import collections
import concurrent.futures
import datetime
import functools
import logging
import multiprocessing
import multiprocessing.pool
import os
import pandas
import pathlib
import queue
import time
import traceback
from typing import Any, Type, Optional, Union, Callable, Sequence  # , Dict, List, Set, Tuple
import vnpy.tools.database_cta
import vnpy.trader.object
import vnpy.trader.utility
import vnpy_ctastrategy.backtesting
import vnpy_ctastrategy.base
from vnpy.trader.constant import Direction, Exchange, Interval, Offset, Status, Product, OptionType, OrderType
from vnpy_zxtools.ctabacktester.cta_config import CtaConfigLite, CtaConfig


class Consumer(object):
    """"""

    def __init__(self, queue_is_none: bool = False) -> None:
        self.queue: Optional[queue.Queue] = None if queue_is_none else multiprocessing.Manager().Queue()
        self.process: multiprocessing.Process = None

    def start(self) -> None:
        """"""
        if (self.queue is not None) and (self.process is None):
            self.process = multiprocessing.Process(target=self.process_entry, args=(self.queue,))
            self.process.start()

    def join(self) -> None:
        """"""
        if self.process is not None:
            self.process.join()

    def put_none(self) -> None:
        """"""
        if self.queue is not None:
            self.queue.put(item=None)

    @classmethod
    def process_entry(cls, que: queue.Queue) -> None:
        """要求该函数能自动结束"""
        while cls.batch(que=que):
            time.sleep(0.1)

    @classmethod
    def batch(cls, que: queue.Queue) -> bool:
        """"""
        is_alive: bool = True

        items: list[vnpy.tools.database_cta.CtaStatisticsData] = []

        while True:
            try:
                item = que.get_nowait()

                if item is None:  # 一旦收到None就认为工作完成了,
                    is_alive: bool = False
                    break

                assert isinstance(item, vnpy.tools.database_cta.CtaStatisticsData)
                items.append(item)

            except queue.Empty as ex:
                break

        if items:
            database: vnpy.tools.database_cta.CtaBacktestDatabase = vnpy.tools.database_cta.get_database()
            database.save_cta_statistics_data(lines=items)

        return is_alive


def calc_once_by_file(
    sqlite_filepath: Union[str, pathlib.Path],
    start: datetime.date,
    end: datetime.date,
    output: bool = False,
    queue: Optional[queue.Queue] = None,
) -> None:
    """"""
    # This function does nothing if the root logger already has handlers configured, unless ...
    logging.basicConfig(level=logging.DEBUG, format="%(message)s")
    logging.getLogger("peewee").setLevel(logging.INFO)

    cta_db: vnpy.tools.database_cta.CtaBacktestDatabase = vnpy.tools.database_cta.get_database()
    cta_db.reinitialize(sqlite_filepath=sqlite_filepath)

    ################################################################################

    cta_dump_data: vnpy.tools.database_cta.CtaDumpData = vnpy.tools.database_cta.CtaDumpDataTool.filepath_2_CtaDumpData(filepath=sqlite_filepath)

    ################################################################################

    logging.info(f"{datetime.datetime.now()}{os.getpid():>9d}\t载入CtaDailyResultData, 开始")
    cta_daily_result_data_list: list[vnpy.tools.database_cta.CtaDailyResultData] = cta_db.load_cta_daily_result_data(
        # 裁剪该数据, 减少空间占用, 所以不能使用全量的查询条件,
        # strategy_name=cta_dump_data.strategy_name,
        # interval=cta_dump_data.interval,
        # exchange=cta_dump_data.exchange,
        # symbol=cta_dump_data.symbol,
        # setting=cta_dump_data.setting,
        date_ge=start,
        date_le=end,
    )
    logging.info(f"{datetime.datetime.now()}{os.getpid():>9d}\t载入CtaDailyResultData, 结束")

    self_daily_results: dict[datetime.date, vnpy_ctastrategy.backtesting.DailyResult] = {}
    for cta_daily_result_data in cta_daily_result_data_list:
        daily_result: vnpy_ctastrategy.backtesting.DailyResult = cta_daily_result_data.py2DailyResult(py_line=cta_daily_result_data)
        self_daily_results[daily_result.date] = daily_result

    self_daily_results = dict(sorted(self_daily_results.items()))

    # Generate dataframe
    results: collections.defaultdict = collections.defaultdict(list)

    for daily_result in self_daily_results.values():
        for key, value in daily_result.__dict__.items():
            if key == "trades":  # 这个字段在DataFrame里面没有用到,
                continue
            results[key].append(value)

    self_daily_df: Optional[pandas.DataFrame] = pandas.DataFrame.from_dict(results).set_index("date") if results else None

    ################################################################################

    if True:
        logging.info(f"{datetime.datetime.now()}{os.getpid():>9d}\t载入CtaStrategyPickleData, 开始")
        cta_strategy_pickle_data_list = cta_db.load_cta_strategy_pickle_data(
            # 不裁剪该数据, 所以可以使用全量的查询条件,
            strategy_name=cta_dump_data.strategy_name,
            interval=cta_dump_data.interval,
            exchange=cta_dump_data.exchange,
            symbol=cta_dump_data.symbol,
            setting=cta_dump_data.setting,
        )
        logging.info(f"{datetime.datetime.now()}{os.getpid():>9d}\t载入CtaStrategyPickleData, 结束")
        assert len(cta_strategy_pickle_data_list) == 1
        cta_strategy_pickle_data: vnpy.tools.database_cta.CtaStrategyPickleData = cta_strategy_pickle_data_list[0]
        assert cta_dump_data.setting == cta_strategy_pickle_data.setting

    vt_symbol: str = f"{cta_strategy_pickle_data.symbol}.{cta_strategy_pickle_data.exchange.value}"

    engine: vnpy_ctastrategy.backtesting.BacktestingEngine = vnpy_ctastrategy.backtesting.BacktestingEngine()

    # TODO: 为什么 print 无法输出到控制台, logging 却可以呢?
    engine.output = lambda msg: logging.info(f"{datetime.datetime.now()}{engine.pid:>9d}\t{msg}")

    engine.set_parameters(
        vt_symbol=vt_symbol,
        interval=cta_strategy_pickle_data.interval,
        start=datetime.datetime.combine(start, datetime.time()),  # 函数 calculate_statistics 没用到它,
        rate=cta_strategy_pickle_data.rate,
        slippage=cta_strategy_pickle_data.slippage,
        size=cta_strategy_pickle_data.size,
        pricetick=cta_strategy_pickle_data.pricetick,
        capital=cta_strategy_pickle_data.capital,  # 函数 calculate_statistics 使用了它,
        end=datetime.datetime.combine(end, datetime.time()),  # 函数 calculate_statistics 没用到它,
        mode=vnpy_ctastrategy.base.BacktestingMode.BAR,  # 函数 calculate_statistics 没用到它,
        risk_free=cta_strategy_pickle_data.risk_free,  # 函数 calculate_statistics 使用了它,
        annual_days=cta_strategy_pickle_data.annual_days,  # 函数 calculate_statistics 使用了它,
        half_life=cta_strategy_pickle_data.half_life,  # 函数 calculate_statistics 使用了它,
        do_continue=None,  # 函数 calculate_statistics 没用到它,
    )

    ################################################################################

    # 拷贝 vnpy_ctastrategy.backtesting.BacktestingEngine.calculate_statistics 函数的实现到文件 file.py 然后执行如下命令:
    # cat file.py | grep -Eo "self\.[a-zA-Z0-9_]+" | sort | uniq | tr "\n" ","
    # 可以知道这个函数使用了这个类的对象的哪些属性,
    statistics: dict = engine.calculate_statistics(df=self_daily_df, output=output)

    ################################################################################

    # 拷贝 vnpy_ctastrategy.backtesting.generate_CtaStatisticsData 函数的实现到文件 file.py 然后执行如下命令:
    # cat file.py | grep -Eo "engine\.[a-zA-Z0-9_]+" | sort | uniq | tr "\n" ","
    # 可以知道这个函数使用了这个类的对象的哪些属性,
    cta_statistics_data = vnpy_ctastrategy.backtesting.generate_CtaStatisticsData(
        engine=engine,
        strategy_name=cta_strategy_pickle_data.strategy_name,
        setting=cta_strategy_pickle_data.setting,
        statistics=statistics,
    )

    if queue is None:
        vnpy_ctastrategy.backtesting.save_cta_backtest_result_bf(cta_statistics_data=cta_statistics_data)
    else:
        queue.put(cta_statistics_data)


def calc_multi_by_file(
    strategy_name: str,
    interval: Interval,
    exchange: Exchange,
    symbol: str,
    start: datetime.date,
    end: datetime.date,
) -> None:
    """"""
    glob_pattern = "{strategy_name}_{interval}_{exchange}_{symbol}_*.db".format(
        strategy_name=strategy_name,
        interval=interval.value,
        exchange=exchange.value,
        symbol=symbol,
    )

    dir_path: pathlib.Path = vnpy.trader.utility.get_folder_path(vnpy.tools.database_cta.CTA_BACKTEST_RESULT_BF_CONT_DBS)

    filepaths: list[pathlib.Path] = [_ for _ in dir_path.glob(glob_pattern) if _.is_file()]

    consumer = Consumer(queue_is_none=False)
    consumer.start()

    func_list: list[callable] = []
    for filepath in filepaths:
        func: callable = functools.partial(
            calc_once_by_file,
            filepath.as_posix(),
            start,
            end,
            False,
            consumer.queue,
        )
        func_list.append(func)

    results_num: int = 0

    percentage: float = results_num / len(func_list) * 100
    logging.info("计算进度：{:d}/{:d}, {:.2f}%".format(results_num, len(func_list), percentage))

    with concurrent.futures.ProcessPoolExecutor(
        max_workers=None,
        mp_context=multiprocessing.get_context("spawn")
    ) as executor:
        fs = {executor.submit(func) for func in func_list}
        for future in concurrent.futures.as_completed(fs):
            assert isinstance(future, concurrent.futures.Future)
            try:
                future.result()  # 如果任务异常了, 会通过它感知到,
            except Exception as exception:
                logging.info(f"exception={exception}; " + traceback.format_exc())

            results_num += 1

            if results_num % 1000 == 0:
                percentage: float = results_num / len(func_list) * 100
                logging.info("计算进度：{:d}/{:d}, {:.2f}%".format(results_num, len(func_list), percentage))

    percentage: float = results_num / len(func_list) * 100
    logging.info("计算进度：{:d}/{:d}, {:.2f}%".format(results_num, len(func_list), percentage))

    consumer.put_none()
    consumer.join()

    if consumer.queue is not None:
        assert consumer.queue.empty()
        logging.info("计算进度：{:d}/{:d}, {:.2f}%, 保存结束".format(results_num, len(func_list), percentage))


def calc_once_by_query(
    strategy_name: str,
    interval: Interval,
    exchange: Exchange,
    symbol: str,
    setting: dict,
    start: datetime.date,
    end: datetime.date,
    cta_strategy_pickle_data: vnpy.tools.database_cta.CtaStrategyPickleData = None,
    output: bool = False,
) -> None:
    """"""
    # This function does nothing if the root logger already has handlers configured, unless ...
    logging.basicConfig(level=logging.DEBUG, format="%(message)s")
    logging.getLogger("peewee").setLevel(logging.INFO)

    cta_db: vnpy.tools.database_cta.CtaBacktestDatabase = vnpy.tools.database_cta.get_database()

    ################################################################################

    logging.info(f"{datetime.datetime.now()}{os.getpid():>9d}\t载入CtaDailyResultData, 开始")
    cta_daily_result_data_list: list[vnpy.tools.database_cta.CtaDailyResultData] = cta_db.load_cta_daily_result_data(
        strategy_name=strategy_name,
        interval=interval,
        exchange=exchange,
        symbol=symbol,
        setting=setting,
        date_ge=start,
        date_le=end,
    )
    logging.info(f"{datetime.datetime.now()}{os.getpid():>9d}\t载入CtaDailyResultData, 结束")

    self_daily_results: dict[datetime.date, vnpy_ctastrategy.backtesting.DailyResult] = {}
    for cta_daily_result_data in cta_daily_result_data_list:
        daily_result: vnpy_ctastrategy.backtesting.DailyResult = cta_daily_result_data.py2DailyResult(py_line=cta_daily_result_data)
        self_daily_results[daily_result.date] = daily_result

    self_daily_results = dict(sorted(self_daily_results.items()))

    # Generate dataframe
    results: collections.defaultdict = collections.defaultdict(list)

    for daily_result in self_daily_results.values():
        for key, value in daily_result.__dict__.items():
            if key == "trades":  # 这个字段在DataFrame里面没有用到,
                continue
            results[key].append(value)

    self_daily_df: Optional[pandas.DataFrame] = pandas.DataFrame.from_dict(results).set_index("date") if results else None

    ################################################################################

    if cta_strategy_pickle_data is None:
        logging.info(f"{datetime.datetime.now()}{os.getpid():>9d}\t载入CtaStrategyPickleData, 开始")
        cta_strategy_pickle_data_list = cta_db.load_cta_strategy_pickle_data(
            strategy_name=strategy_name,
            interval=interval,
            exchange=exchange,
            symbol=symbol,
            setting=setting,
        )
        logging.info(f"{datetime.datetime.now()}{os.getpid():>9d}\t载入CtaStrategyPickleData, 结束")
        assert len(cta_strategy_pickle_data_list) == 1
        cta_strategy_pickle_data: vnpy.tools.database_cta.CtaStrategyPickleData = cta_strategy_pickle_data_list[0]

    vt_symbol: str = f"{cta_strategy_pickle_data.symbol}.{cta_strategy_pickle_data.exchange.value}"

    engine: vnpy_ctastrategy.backtesting.BacktestingEngine = vnpy_ctastrategy.backtesting.BacktestingEngine()

    # TODO: 为什么 print 无法输出到控制台, logging 却可以呢?
    engine.output = lambda msg: logging.info(f"{datetime.datetime.now()}{engine.pid:>9d}\t{msg}")

    engine.set_parameters(
        vt_symbol=vt_symbol,
        interval=cta_strategy_pickle_data.interval,
        start=datetime.datetime.combine(start, datetime.time()),  # 函数 calculate_statistics 没用到它,
        rate=cta_strategy_pickle_data.rate,
        slippage=cta_strategy_pickle_data.slippage,
        size=cta_strategy_pickle_data.size,
        pricetick=cta_strategy_pickle_data.pricetick,
        capital=cta_strategy_pickle_data.capital,  # 函数 calculate_statistics 使用了它,
        end=datetime.datetime.combine(end, datetime.time()),  # 函数 calculate_statistics 没用到它,
        mode=vnpy_ctastrategy.base.BacktestingMode.BAR,  # 函数 calculate_statistics 没用到它,
        risk_free=cta_strategy_pickle_data.risk_free,  # 函数 calculate_statistics 使用了它,
        annual_days=cta_strategy_pickle_data.annual_days,  # 函数 calculate_statistics 使用了它,
        half_life=cta_strategy_pickle_data.half_life,  # 函数 calculate_statistics 使用了它,
        do_continue=None,  # 函数 calculate_statistics 没用到它,
    )

    ################################################################################

    # 拷贝 vnpy_ctastrategy.backtesting.BacktestingEngine.calculate_statistics 函数的实现到文件 file.py 然后执行如下命令:
    # cat file.py | grep -Eo "self\.[a-zA-Z0-9_]+" | sort | uniq | tr "\n" ","
    # 可以知道这个函数使用了这个类的对象的哪些属性,
    statistics: dict = engine.calculate_statistics(df=self_daily_df, output=output)

    ################################################################################

    # 拷贝 vnpy_ctastrategy.backtesting.generate_CtaStatisticsData 函数的实现到文件 file.py 然后执行如下命令:
    # cat file.py | grep -Eo "engine\.[a-zA-Z0-9_]+" | sort | uniq | tr "\n" ","
    # 可以知道这个函数使用了这个类的对象的哪些属性,
    cta_statistics_data = vnpy_ctastrategy.backtesting.generate_CtaStatisticsData(
        engine=engine,
        strategy_name=cta_strategy_pickle_data.strategy_name,
        setting=cta_strategy_pickle_data.setting,
        statistics=statistics,
    )

    vnpy_ctastrategy.backtesting.save_cta_backtest_result_bf(cta_statistics_data=cta_statistics_data)


def calc_multi_by_query(
    strategy_name: str,
    interval: Interval,
    exchange: Exchange,
    symbol: str,
    start: datetime.date,
    end: datetime.date,
) -> None:
    """"""
    cta_db: vnpy.tools.database_cta.CtaBacktestDatabase = vnpy.tools.database_cta.get_database()

    cta_strategy_pickle_data_list = cta_db.load_cta_strategy_pickle_data(
        strategy_name=strategy_name,
        interval=interval,
        exchange=exchange,
        symbol=symbol,
        setting=None,
    )

    func_list: list[callable] = []
    for cta_strategy_pickle_data in cta_strategy_pickle_data_list:
        assert strategy_name == cta_strategy_pickle_data.strategy_name
        assert interval == cta_strategy_pickle_data.interval
        assert exchange == cta_strategy_pickle_data.exchange
        assert symbol == cta_strategy_pickle_data.symbol
        func: callable = functools.partial(
            calc_once_by_query,
            cta_strategy_pickle_data.strategy_name,
            cta_strategy_pickle_data.interval,
            cta_strategy_pickle_data.exchange,
            cta_strategy_pickle_data.symbol,
            cta_strategy_pickle_data.setting,
            start,
            end,
            cta_strategy_pickle_data,
            False,
        )
        func_list.append(func)

    results_num: int = 0

    percentage: float = results_num / len(func_list) * 100
    logging.info("计算进度：{:d}/{:d}, {:.2f}%".format(results_num, len(func_list), percentage))

    with concurrent.futures.ProcessPoolExecutor(
        max_workers=None,
        mp_context=multiprocessing.get_context("spawn")
    ) as executor:
        fs = {executor.submit(func) for func in func_list}
        for future in concurrent.futures.as_completed(fs):
            assert isinstance(future, concurrent.futures.Future)
            try:
                future.result()  # 如果任务异常了, 会通过它感知到,
            except Exception as exception:
                logging.info(f"exception={exception}; " + traceback.format_exc())

            results_num += 1

            if results_num % 1000 == 0:
                percentage: float = results_num / len(func_list) * 100
                logging.info("计算进度：{:d}/{:d}, {:.2f}%".format(results_num, len(func_list), percentage))

    percentage: float = results_num / len(func_list) * 100
    logging.info("计算进度：{:d}/{:d}, {:.2f}%".format(results_num, len(func_list), percentage))


def initialize_logging():
    """"""
    c_handler: logging.StreamHandler = logging.StreamHandler()
    f_handler: logging.FileHandler = logging.FileHandler(filename=f"{__file__}.log", mode="a", encoding="utf8")
    logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(levelname)s - %(message)s", handlers=[c_handler, f_handler])
    logging.getLogger("peewee").setLevel(logging.INFO)


if __name__ == '__main__':
    initialize_logging()

    logging.info(f"++++++++++ START ++++++++++")

    parser = argparse.ArgumentParser()
    parser.add_argument("--file", "-f", type=str, default="", help="json配置文件")
    # parser.add_argument 中的 type 参数可以是任何接受单个字符串的可调用参数, 此时, argparse 只是将该字符串传递给 type 的构造函数, 因此可以如下书写:
    parser.add_argument("--beg", "-b", type=lambda d: datetime.datetime.strptime(d, "%Y-%m-%d"), default=None, help="开始日期(begin)")
    parser.add_argument("--end", "-e", type=lambda d: datetime.datetime.strptime(d, "%Y-%m-%d"), default=None, help="结束日期(end)")
    args = parser.parse_args()

    json_file: pathlib.Path = pathlib.Path(args.file)

    if not json_file.is_file():
        logging.info(f"json_file={json_file.absolute()}, not is_file, will exit...")
        exit(1)

    logging.info(f"json_file={json_file.absolute()}, beg={args.beg}, end={args.end},")

    cta_config_lite: CtaConfigLite = CtaConfigLite.load_from_file(filepath=json_file)

    if args.beg:
        assert isinstance(args.beg, type(cta_config_lite.start))
        cta_config_lite.start = args.beg
    if args.end:
        assert isinstance(args.end, type(cta_config_lite.end))
        cta_config_lite.end = args.end

    logging.info(cta_config_lite)

    symbol, exchange = vnpy.trader.utility.extract_vt_symbol(vt_symbol=cta_config_lite.vt_symbol)

    beg = datetime.datetime.now()
    calc_multi_by_file(
        strategy_name=cta_config_lite.strategy_class.__name__,
        interval=cta_config_lite.interval,
        exchange=exchange,
        symbol=symbol,
        start=cta_config_lite.start.date(),
        end=cta_config_lite.end.date(),
    )
    end = datetime.datetime.now()

    cost = (end - beg).total_seconds()
    logging.info(f"耗时{cost}秒")
